#!/usr/bin/env python
import freenect
import opencv.cv as cv
import numpy as np
import time
import cPickle
import sys

cv.NamedWindow('Depth')
cv.NamedWindow('RGB')

depth_buffer = []
image_buffer = []
name = sys.argv[1]

def display_depth(dev, data, timestamp):
    if time.time() - start_time > 3.:
        depth_buffer.append(data)
    data -= np.min(data.ravel())
    data *= 65536 / np.max(data.ravel())
    image = cv.CreateImageHeader((data.shape[1], data.shape[0]),
                                 cv.IPL_DEPTH_16U,
                                 1)
    cv.SetData(image, data.tostring(),
               data.dtype.itemsize * data.shape[1])
    cv.ShowImage('Depth', image)
    cv.WaitKey(5)


def display_rgb(dev, data, timestamp):
    if time.time() - start_time > 3.:
        image_buffer.append(data)
    image = cv.CreateImageHeader((data.shape[1], data.shape[0]),
                                 cv.IPL_DEPTH_8U,
                                 3)
    # Note: We swap from RGB to BGR here
    cv.SetData(image, data[:, :, ::-1].tostring(),
               data.dtype.itemsize * 3 * data.shape[1])
    cv.ShowImage('RGB', image)
    cv.WaitKey(5)

def body(dev, ctx):
    if time.time() - start_time > 8.:
        depth_file = open(name + '_depth.data', 'w')
        image_file = open(name + '_image.data', 'w')

        cPickle.dump(depth_buffer,depth_file)
        cPickle.dump(image_buffer,image_file)
        print('10 sec passed, killing')
        raise freenect.Kill

start_time = time.time()
freenect.runloop(depth=display_depth,
                 video=display_rgb, body=body)
